import pandas as pd
from sqlalchemy import func, and_
from datetime import datetime, timedelta
from app import db
from app.models.order import Order, OrderItem
from app.models.blind_box import BlindBox

class DataLoader:
    @staticmethod
    def get_sales_data(days=365, category=None, box_id=None):
        """获取销售数据，默认获取过去一年的数据"""
        end_date = datetime.now()
        start_date = end_date - timedelta(days=days)
        
        # 构建基本查询
        query = db.session.query(
            func.date(Order.create_time).label('date'),
            BlindBox.category,
            BlindBox.box_id,
            BlindBox.name,
            func.sum(OrderItem.quantity).label('sales_quantity'),
            func.sum(OrderItem.price * OrderItem.quantity).label('sales_amount'),
            func.avg(OrderItem.price).label('avg_price')
        ).join(
            OrderItem, Order.order_id == OrderItem.order_id
        ).join(
            BlindBox, OrderItem.box_id == BlindBox.box_id
        ).filter(
            Order.create_time.between(start_date, end_date),
            Order.status.in_([2, 3])  # 已支付和已完成的订单
        )
        
        # 应用过滤条件
        if category:
            query = query.filter(BlindBox.category == category)
        if box_id:
            query = query.filter(BlindBox.box_id == box_id)
        
        # 按日期、类别和盲盒ID分组
        query = query.group_by(
            func.date(Order.create_time),
            BlindBox.category,
            BlindBox.box_id,
            BlindBox.name
        ).order_by(
            func.date(Order.create_time)
        )
        
        # 执行查询并转换为DataFrame
        result = query.all()
        columns = ['date', 'category', 'box_id', 'name', 'sales_quantity', 'sales_amount', 'avg_price']
        df = pd.DataFrame(result, columns=columns)
        
        # 转换数据类型
        df['date'] = pd.to_datetime(df['date'])
        df['sales_quantity'] = df['sales_quantity'].astype(int)
        df['sales_amount'] = df['sales_amount'].astype(float)
        df['avg_price'] = df['avg_price'].astype(float)
        
        return df
    
    @staticmethod
    def prepare_time_series_data(df, freq='D', fill_missing=True):
        """准备时间序列数据，包括缺失值处理和重采样"""
        # 设置日期为索引
        df = df.set_index('date')
        
        if fill_missing:
            # 获取日期范围
            min_date = df.index.min()
            max_date = df.index.max()
            
            # 创建完整日期范围
            date_range = pd.date_range(start=min_date, end=max_date, freq=freq)
            
            # 重新索引并填充缺失值
            df = df.reindex(date_range)
            df['sales_quantity'] = df['sales_quantity'].fillna(0)
            
            # 如果存在sales_amount列，则填充
            if 'sales_amount' in df.columns:
                df['sales_amount'] = df['sales_amount'].fillna(0)
            
            # 仅当相应列存在时填充
            if 'category' in df.columns:
                df['category'] = df['category'].ffill().bfill()
            
            if 'box_id' in df.columns:
                df['box_id'] = df['box_id'].ffill().bfill()
            
            if 'name' in df.columns:
                df['name'] = df['name'].ffill().bfill()
            
            # 对价格使用平均值填充，如果存在
            if 'avg_price' in df.columns:
                # 确保有数据才计算均值
                mean_price = df['avg_price'].mean()
                if pd.notna(mean_price):
                    df['avg_price'] = df['avg_price'].fillna(mean_price)
                else:
                    df['avg_price'] = df['avg_price'].fillna(0)
        
        return df
    
    @staticmethod
    def aggregate_by_category(df, freq='D'):
        """按类别聚合销售数据"""
        # 确保日期是索引
        if 'date' in df.columns:
            df = df.set_index('date')
        
        # 按类别和日期分组并聚合
        grouped = df.groupby(['category', pd.Grouper(freq=freq)])
        
        # 计算聚合指标
        agg_df = grouped.agg({
            'sales_quantity': 'sum',
            'sales_amount': 'sum',
            'avg_price': 'mean'
        }).reset_index()
        
        return agg_df
